package com.liu.audio.tts;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.audio.tts.SpeechSynthesisResult;
import com.alibaba.dashscope.audio.ttsv2.SpeechSynthesisAudioFormat;
import com.alibaba.dashscope.audio.ttsv2.SpeechSynthesisParam;
import com.alibaba.dashscope.audio.ttsv2.SpeechSynthesizer;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;

import javax.sound.sampled.*;
import java.io.BufferedOutputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;

public class CosyVoiceStream {
    private static String apikey = "sk-d1c41de66c074fc686488ab227d32fbe";
    private static String model = "cosyvoice-v1";
    private static String voice = "longxiaochun";
    public static void process() throws NoApiKeyException, InputRequiredException {
        PlaybackRunnable playbackRunnable = new PlaybackRunnable();
        try {
            playbackRunnable.prepare();
        } catch (LineUnavailableException e) {
            throw new RuntimeException(e);
        }
        Thread playbackThread = new Thread(playbackRunnable);
        // Start the playback thread
        playbackThread.start();


        /*******  Call the Generative AI Model to get streaming text *******/
        // Prepare for the LLM call
        Generation gen = new Generation();
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content("请介绍一下你自己")
                .build();
        GenerationParam genParam = GenerationParam.builder()
                        // 若没有将API Key配置到环境变量中，需将下面这行代码注释放开，并将apiKey替换为自己的API Key
                         .apiKey(apikey)
                        .model("qwen-turbo")
                        .messages(Arrays.asList(userMsg))
                        .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                        .topP(0.8)
                        .incrementalOutput(true)
                        .build();
        Flowable<GenerationResult> result = gen.streamCall(genParam);

        // Prepare the speech synthesis task
        SpeechSynthesisParam param = SpeechSynthesisParam.builder()
                        // 若没有将API Key配置到环境变量中，需将下面这行代码注释放开，并将apiKey替换为自己的API Key
                         .apiKey(apikey)
                        .model(model)
                        .voice(voice)
                        .format(SpeechSynthesisAudioFormat.PCM_22050HZ_MONO_16BIT)
                        .build();
        SpeechSynthesizer synthesizer = new SpeechSynthesizer(param, new ReactCallback(playbackRunnable));


        StringBuilder stringBuilder = new StringBuilder();
        result.blockingForEach(message -> {
            String text = message.getOutput().getChoices().get(0).getMessage().getContent();
            stringBuilder.append(text);
            if (!"".equals(text)){
                System.out.println("LLM output：" + text);
                synthesizer.streamingCall(text);
            }
        });

        synthesizer.streamingComplete();
        System.out.print("requestId: " + synthesizer.getLastRequestId());
        System.out.println("res: " +stringBuilder);
        try {
            // Wait for the playback thread to finish playing all
            playbackThread.join();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    public static void main(String[] args) throws NoApiKeyException, InputRequiredException {
        process();
        System.exit(0);
    }



    // Playback thread 播放线程
    static class PlaybackRunnable implements Runnable {
        // Set the audio format. Please configure according to your actual device,
        // synthesized audio parameters, and platform choice Here it is set to
        // 22050Hz16bit single channel. It is recommended that customers choose other
        // sample rates and formats based on the model sample rate and device
        // compatibility.

        // 设置音频格式，请根据实际自身设备，合成音频参数和平台选择配置
        // 这里选择48k16bit单通道，建议客户根据选用的模型采样率情况和自身设备兼容性选择其他采样率和格式
        private AudioFormat af = new AudioFormat(22050, 16, 1, true, false);
        private DataLine.Info info = new DataLine.Info(SourceDataLine.class, af);
        private SourceDataLine targetSource = null;
        private AtomicBoolean runFlag = new AtomicBoolean(true);
        private ConcurrentLinkedQueue<ByteBuffer> queue = new ConcurrentLinkedQueue<>();


        // Prepare the player
        public void prepare() throws LineUnavailableException {

            targetSource = (SourceDataLine) AudioSystem.getLine(info);
            targetSource.open(af, 4096);
            targetSource.start();
        }

        public void put(ByteBuffer buffer) {

            queue.add(buffer);
        }

        // Stop playback
        public void stop() {
            runFlag.set(false);
        }

        @Override
        public void run() {
            if (targetSource == null) {
                return;
            }

            while (runFlag.get()) {
                if (queue.isEmpty()) {
                    try {
                        Thread.sleep(100);
                    } catch (InterruptedException e) {
                    }
                    continue;
                }

                ByteBuffer buffer = queue.poll();
                if (buffer == null) {
                    continue;
                }

                byte[] data = buffer.array();
                targetSource.write(data, 0, data.length);
            }

            // Play all remaining cache
            if (!queue.isEmpty()) {
                ByteBuffer buffer = null;
                while ((buffer = queue.poll()) != null) {
                    byte[] data = buffer.array();
                    targetSource.write(data, 0, data.length);
                }
            }
            // Release the player
            targetSource.drain();
            targetSource.stop();
            targetSource.close();
        }
    }

    // Create a subclass inheriting from ResultCallback<SpeechSynthesisResult>
    // to implement the callback interface
    // 创建一个继承自ResultCallback<SpeechSynthesisResult>的子类来实现回调接口
    static class ReactCallback extends ResultCallback<SpeechSynthesisResult> {
        private PlaybackRunnable playbackRunnable = null;
        BufferedOutputStream out ;
        FileOutputStream fos;

        public ReactCallback(PlaybackRunnable playbackRunnable) {
            try {
                fos = new FileOutputStream("D:\\output.wav");
            } catch (FileNotFoundException e) {
                throw new RuntimeException(e);
            }
            this.playbackRunnable = playbackRunnable;
        }

        // Callback when the service side returns the streaming synthesis result
        // 当服务侧返回流式合成结果后回调
        @Override
        public void onEvent(SpeechSynthesisResult result) {
            // Get the binary data of the streaming result via getAudio
            // 通过getAudio获取流式结果二进制数据
            if (result.getAudioFrame() != null) {
                ByteBuffer audioFrame = result.getAudioFrame();
                try {
                    fos.write(audioFrame.array());
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }

                // Stream the data to the player
                // 将数据流式推给播放器
                playbackRunnable.put(result.getAudioFrame());
            }
        }

        // Callback when the service side completes the synthesis
        // 当服务侧完成合成后回调
        @Override
        public void onComplete() {
            // Notify the playback thread to end
            try {
                fos.close();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            // 告知播放线程结束
            playbackRunnable.stop();
        }

        // Callback when an error occurs
        // 当出现错误时回调
        @Override
        public void onError(Exception e) {
            // Tell the playback thread to end
            System.out.println(e);
            playbackRunnable.stop();
        }
    }
}